from typing import List
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class TutorialSearch(BaseModel):
    origin: str = Field(description = "用户输入的原始问题")
    domain: str = Field(description = "用户输入问题归属的领域")
    keyword: List[str] = Field(description = "用户输入问题中的关键字")
    purpose: str = Field(description = "用户问题的意图或目的")
    date: str = Field(description = "用户输入问题中的日期")
    address: str = Field(description = "用户输入问题中的地址")
    person: str = Field(description = "用户输入问题中的人物")

    def pretty_search_text(self):
        text = ""
        for field in self.model_fields:
            if getattr(self, field) is not None and getattr(self, field) != getattr(self.model_fields[field], "default", None):
                text += f"\n{field}: {getattr(self, field)}"
        return text

class QueryConstruction:

    def __init__(self, llm_model, vector_store):
        self.llm_model = llm_model
        self.vector_store = vector_store

    def invoke(self, query):
        logger.info(f"QueryConstruction invoke query: {query}")
        retriever = self.vector_store.as_retriever(search_kwargs={"k": 3})

        tutorial_template = """
            你是一个问题分析转换大师，分析用户输入的问题，并分析转化。
            用户问题：{question}
        """
        tutorial_prompt = ChatPromptTemplate.from_template(tutorial_template)
        tutorial_chain = tutorial_prompt | self.llm_model.with_structured_output(TutorialSearch)
        tutorial_result = tutorial_chain.invoke({ "question": query })
        logger.info(f"QueryConstruction invoke tutorial_result: {tutorial_result}")

        content = tutorial_result.pretty_search_text()
        logger.info(f"QueryConstruction invoke content: {content}")
        retrieve_docs = retriever.invoke(content)
        retrieve_doc = "\n".join([doc.page_content for doc in retrieve_docs])
        logger.info(f"QueryConstruction invoke retrieve_doc len: {len(retrieve_doc)}")

        template = """
            请基于以下上下文内容回答问题：
            {context}
            
            问题：{question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        chain_result = chain.invoke({ "context": retrieve_doc, "question": query })
        logger.info(f"QueryConstruction invoke chain_result len: {len(chain_result)}")
        return { "retrieve_docs": retrieve_doc, "chain_result": chain_result }
